from flask import Flask, request, jsonify, make_response, redirect, render_template
import jwt
import datetime
from functools import wraps

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-very-secret-key-123!@#'  # 生产环境应使用更安全的密钥
app.config['TOKEN_EXPIRATION'] = 30  # Token有效期(分钟)

# 模拟用户数据库
users = {
    'admin': {'password': 'adminpass', 'role': 'admin', 'name': 'Administrator'},
    'alice': {'password': 'alicepass', 'role': 'editor', 'name': 'Alice Smith'},
    'bob': {'password': 'bobpass', 'role': 'viewer', 'name': 'Bob Johnson'}
}


# 登录页面
@app.route('/')
def index():
    return render_template('login.html')


# 登录接口 - 生成Token
@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    if not username or not password:
        return jsonify({'error': 'Username and password required'}), 400

    user = users.get(username)

    if not user or user['password'] != password:
        return jsonify({'error': 'Invalid credentials'}), 401

    # 生成JWT Token
    token = jwt.encode({
        'sub': username,
        'name': user['name'],
        'role': user['role'],
        'iat': datetime.datetime.utcnow(),
        'exp': datetime.datetime.utcnow() + datetime.timedelta(minutes=app.config['TOKEN_EXPIRATION'])
    }, app.config['SECRET_KEY'], algorithm='HS256')

    # 创建响应 - 可以选择通过Cookie或响应体返回Token
    response = jsonify({
        'message': 'Login successful',
        'user': {
            'username': username,
            'name': user['name'],
            'role': user['role']
        }
    })

    # 通过Cookie设置Token
    response.set_cookie(
        'access_token',
        token,
        httponly=True,  # 防止XSS攻击
        secure=True,  # 仅通过HTTPS传输
        samesite='Strict',  # 防止CSRF攻击
        max_age=app.config['TOKEN_EXPIRATION'] * 60  # 过期时间(秒)
    )

    return response


# Token验证装饰器
def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = None

        # 1. 首先尝试从Cookie获取Token
        token = request.cookies.get('access_token')

        # 2. 如果Cookie中没有，尝试从Authorization头获取
        if not token:
            auth_header = request.headers.get('Authorization')
            if auth_header and auth_header.startswith('Bearer '):
                token = auth_header.split(' ')[1]

        # 3. 如果都没有，尝试从自定义头获取
        if not token:
            token = request.headers.get('X-Access-Token')

        if not token:
            return jsonify({'error': 'Token is missing!'}), 401

        try:
            # 解码并验证Token
            data = jwt.decode(token, app.config['SECRET_KEY'], algorithms=['HS256'])
            current_user = {
                'username': data['sub'],
                'name': data['name'],
                'role': data['role']
            }
        except jwt.ExpiredSignatureError:
            return jsonify({'error': 'Token has expired!'}), 401
        except jwt.InvalidTokenError:
            return jsonify({'error': 'Invalid token!'}), 401

        # 将用户信息添加到请求上下文中
        return f(current_user, *args, **kwargs)

    return decorated


# 受保护的路由 - 需要有效Token
@app.route('/protected')
@token_required
def protected(current_user):
    return jsonify({
        'message': f'Hello {current_user["name"]}!',
        'user': current_user,
        'status': 'You have accessed a protected resource'
    })


# 管理员路由 - 需要管理员权限
@app.route('/admin')
@token_required
def admin_panel(current_user):
    if current_user['role'] != 'admin':
        return jsonify({'error': 'Admin access required'}), 403

    return jsonify({
        'message': f'Welcome, Admin {current_user["name"]}!',
        'admin_data': 'Sensitive admin information here'
    })


# Token刷新接口
@app.route('/refresh', methods=['POST'])
@token_required
def refresh_token(current_user):
    # 生成新Token
    user = users.get(current_user['username'])
    new_token = jwt.encode({
        'sub': current_user['username'],
        'name': user['name'],
        'role': user['role'],
        'iat': datetime.datetime.utcnow(),
        'exp': datetime.datetime.utcnow() + datetime.timedelta(minutes=app.config['TOKEN_EXPIRATION'])
    }, app.config['SECRET_KEY'], algorithm='HS256')

    response = jsonify({'message': 'Token refreshed'})
    response.set_cookie('access_token', new_token,
                        httponly=True, secure=True,
                        samesite='Strict',
                        max_age=app.config['TOKEN_EXPIRATION'] * 60)
    return response


# 登出接口
@app.route('/logout', methods=['POST'])
def logout():
    response = jsonify({'message': 'Successfully logged out'})
    # 通过设置过期时间为0来删除Cookie
    response.set_cookie('access_token', '', expires=0)
    return response


if __name__ == '__main__':
    app.run(debug=True, port=5000)